package main

import (
	"fmt"
	"log"
	"math/rand/v2"
	"testing"
)

// CHI is an interface ConsistentHasher types implement, for common testing
type CHI interface {
	FindNodeFor(string) string
	AddNode(string) error
	RemoveNode(string) error
}

func makeLoggedRand(t *testing.T) *rand.Rand {
	s1, s2 := rand.Uint64(), rand.Uint64()
	log.Printf("%s seed: %v, %v", t.Name(), s1, s2)
	return rand.New(rand.NewPCG(s1, s2))
}

func TestSmoke(t *testing.T) {
	t.Run("smokeCH", func(t *testing.T) {
		ch := NewConsistentHasher(10000)
		doTestSmoke(t, ch)
	})
	t.Run("smokeCHV", func(t *testing.T) {
		ch := NewConsistentHasherV(10000)
		doTestSmoke(t, ch)
	})
}

func doTestSmoke(t *testing.T, ch CHI) {
	var nodes []string
	for i := range 8 {
		n := fmt.Sprintf("node-%03d", i)
		nodes = append(nodes, n)
		if err := ch.AddNode(n); err != nil {
			t.Error(err)
		}
	}

	s := "blobs"
	ch.FindNodeFor(s)
}

func TestFindsCorrectNode(t *testing.T) {
	rnd := makeLoggedRand(t)
	ch := NewConsistentHasher(128 * 1024 * 1024)

	// Add nodes named "node-N"
	var nodes []string
	for i := range 256 {
		n := fmt.Sprintf("node-%03d", i)
		nodes = append(nodes, n)
		if err := ch.AddNode(n); err != nil {
			t.Error(err)
		}
	}

	// Repeat many times:
	// - generate a random item
	// - ask the ConsistentHasher which node it hashes to
	// - run a manual search process using ConsistentHasher's internals to verify
	//   that the item was hashed to the right node
	for range 1000 {
		str := generateRandomString(rnd, 16)
		sh := hashItem(str, ch.ringSize)

		nn := ch.FindNodeFor(str)

		// Now find the closest (from above) node in the ConsistentHasher,
		// and verify that's what we got.
		bestDiff := ch.ringSize
		bestSlot := 0
		for i, s := range ch.slots {
			if s >= sh && (s-sh) < bestDiff {
				bestDiff = s - sh
				bestSlot = i
			}
		}

		// Special case if wrap around
		if sh > ch.slots[len(ch.slots)-1] {
			diff := ch.ringSize - sh + ch.slots[0]
			if diff < bestDiff {
				bestDiff = diff
				bestSlot = 0
			}
		}

		if nn != ch.nodes[bestSlot] {
			t.Errorf("mismatch; ch returned %v, manual search %v", nn, ch.nodes[bestSlot])
		}
	}
}

// Similar to TestFindsCorrectNode, but for the V variant - with similar
// checking of the internal slots slice
func TestFindsCorrectNodeV(t *testing.T) {
	rnd := makeLoggedRand(t)
	ch := NewConsistentHasherV(100000000)

	// Add nodes named "node-N"
	var nodes []string
	for i := range 256 {
		n := fmt.Sprintf("node-%03d", i)
		nodes = append(nodes, n)
		if err := ch.AddNode(n); err != nil {
			t.Error(err)
		}
	}

	// Repeat many times:
	// - generate a random item
	// - ask the ConsistentHasherV which node it hashes to
	// - run a manual search process using ConsistentHasherV's internals to verify
	//   that the item was hashed to the right node
	for range 1000 {
		str := generateRandomString(rnd, 16)
		sh := hashItem(str, ch.ringSize)

		nn := ch.FindNodeFor(str)

		// Now find the closest (from above) node in the ConsistentHasherV,
		// and verify that's what we got.
		bestDiff := ch.ringSize
		bestSlot := 0
		for i, s := range ch.slots {
			if s >= sh && (s-sh) < bestDiff {
				bestDiff = s - sh
				bestSlot = i
			}
		}

		// Special case if wrap around
		if sh > ch.slots[len(ch.slots)-1] {
			diff := ch.ringSize - sh + ch.slots[0]
			if diff < bestDiff {
				bestDiff = diff
				bestSlot = 0
			}
		}

		vnode := ch.nodes[bestSlot]
		node := nodeFromVnode(vnode)

		if nn != node {
			t.Errorf("mismatch; ch returned %v, manual search %v", nn, node)
		}
	}
}

func TestConsistentAfterAdd(t *testing.T) {
	t.Run("consistentCH", func(t *testing.T) {
		ch := NewConsistentHasher(1000000)
		doTestConsistentAfterAdd(t, ch)
	})
	t.Run("consistentCHV", func(t *testing.T) {
		ch := NewConsistentHasherV(1000000)
		doTestConsistentAfterAdd(t, ch)
	})
}

func doTestConsistentAfterAdd(t *testing.T, ch CHI) {
	rnd := makeLoggedRand(t)

	// Add nodes named "node-N"
	var nodes []string
	for i := range 32 {
		n := fmt.Sprintf("node-%03d", i)
		nodes = append(nodes, n)
		if err := ch.AddNode(n); err != nil {
			t.Error(err)
		}
	}

	// Generate random items and write down which nodes they hashed to
	mapBeforeAdd := make(map[string]string)
	for range 100000 {
		str := generateRandomString(rnd, 16)
		nn := ch.FindNodeFor(str)
		mapBeforeAdd[str] = nn
	}

	// Now add a new node.
	newNode := "anewnode"
	if err := ch.AddNode(newNode); err != nil {
		t.Error(err)
	}

	// Hash the same items again; the node they are hashed to shouldn't have
	// changed, unless it's to newNode. Also, expect that at least some items
	// were rehashed.
	rehashedCount := 0
	for item, n := range mapBeforeAdd {
		nn := ch.FindNodeFor(item)
		if nn != n {
			if nn != newNode {
				t.Errorf("%v rehashed from %v to %v", item, n, nn)
			}
			rehashedCount += 1
		}
	}
	if rehashedCount == 0 {
		t.Errorf("got rehashedCount=0")
	}
}

func TestRemove(t *testing.T) {
	var size uint64 = 100000000
	t.Run("removeCH", func(t *testing.T) {
		ch := NewConsistentHasher(size)
		doTestRemoveWorks(t, ch)

		ch = NewConsistentHasher(size)
		doTestConsistentAfterRemove(t, ch)
	})
	t.Run("removeCHV", func(t *testing.T) {
		ch := NewConsistentHasherV(size)
		doTestRemoveWorks(t, ch)

		ch = NewConsistentHasherV(size)
		doTestConsistentAfterRemove(t, ch)
	})
}

func doTestRemoveWorks(t *testing.T, ch CHI) {
	// Add nodes named "node-N"
	var nodes []string
	for i := range 256 {
		n := fmt.Sprintf("node-%03d", i)
		nodes = append(nodes, n)
		if err := ch.AddNode(n); err != nil {
			t.Error(err)
		}
	}

	// Nodes that don't exist
	for _, n := range []string{"node-300", "anode", "node-0000"} {
		err := ch.RemoveNode(n)
		if err == nil {
			t.Errorf("got no error removing non-existing node")
		}
	}

	// Remove all existing nodes
	for _, n := range nodes {
		if err := ch.RemoveNode(n); err != nil {
			t.Error(err)
		}
	}
}

func doTestConsistentAfterRemove(t *testing.T, ch CHI) {
	rnd := makeLoggedRand(t)

	// Add nodes named "node-N"
	var nodes []string
	for i := range 32 {
		n := fmt.Sprintf("node-%03d", i)
		nodes = append(nodes, n)
		if err := ch.AddNode(n); err != nil {
			t.Error(err)
		}
	}

	// Generate random items and write down which nodes they hashed to
	mapBefore := make(map[string]string)
	for range 100000 {
		str := generateRandomString(rnd, 16)
		nn := ch.FindNodeFor(str)
		mapBefore[str] = nn
	}

	// Now remove a node
	removedNode := "node-009"
	if err := ch.RemoveNode(removedNode); err != nil {
		t.Error(err)
	}

	// Hash the same items again; the node they are hashed to shouldn't have
	// changed, unless it was removedNode. Also, expect that at least some items
	// were rehashed.
	rehashedCount := 0
	for item, n := range mapBefore {
		nn := ch.FindNodeFor(item)
		if nn != n {
			if n != removedNode {
				t.Errorf("%v rehashed from %v to %v", item, n, nn)
			}
			rehashedCount += 1
		}
	}
	if rehashedCount == 0 {
		t.Errorf("got rehashedCount=0")
	}
}
